'''这个文件为对ann_path_origin和img_path_origin选取特定类别的图片和xml
生成的结果为一个图，一个txt文件,这张图片没有包含目标不会被写入txt
这个会复制提取出来的图片'''

import os
from tqdm import tqdm
import shutil
import xml.etree.ElementTree as ET

ann_path_origin = r'/home/xys/CloundShiProjects/traffic_light/trafficlight_dect/初data/xml/'        #要提取的ann标签文件路径
img_path_origin = r'/home/xys/CloundShiProjects/traffic_light/trafficlight_dect/初data/mark/'         #要提取的Image文件路径
img_savepath = r'/home/xys/CloundShiProjects/traffic_light/trafficlight_dect/data/JPEGImages/'         #要保存到的Images文件路径
txt_save_path = r'/home/xys/datasets/Vehicle/voc/新建文件夹/'   #存入txt的路径
class_names_path=r'/home/xys/datasets/Vehicle/classes.txt'     #类别文件路径

#-----------------------函数定义------------------------------------

def mkr(path):
    '''
    如果path存在，就先删除所有内容再创建，否则直接创建
    '''
    if os.path.exists(path):#先删除再创建
        shutil.rmtree(path)#递归删除目录树
        os.mkdir(path)#创建目录
    else:
        os.mkdir(path)#如果不存在则直接创建

def read_class_name(path):        #读取path下的类别民
    f = open(path,'r')
    classes_name = []
    for i in f.readlines():
        classes_name.append(i.strip())
    return classes_name


names = locals()
classes_name = read_class_name(class_names_path)      #读取类别信息
mkr(img_savepath)                               #创建文件夹
mkr(txt_save_path)

count = 0
for xml_file_name in tqdm(os.listdir(ann_path_origin),ncols=150):
    file_name = xml_file_name[:-4]
    xml_file = open(ann_path_origin+xml_file_name)      #打开xml文件
    tree = ET.parse(xml_file)
    root = tree.getroot()                               #得到根
    image_width = int(root.find('size').find('width').text)
    image_height = int(root.find('size').find('height').text)
    have = False                                        #一个用来判断是否含有该类的标志
    all_lines = []
    for obj in root.iter('object'):
        diffcult = obj.find('difficult').text         #找到difficult，难以识别的类
        cls_name = obj.find('name').text
        if cls_name not in classes_name or int(diffcult) == 1:
            continue

        #--------------------下面的语句表示存在这个boundingbox，将此boundingbox的坐标类别写入到xml里-------------------
        if cls_name in classes_name:
            have = True
        xmlbox = obj.find('bndbox')                 #找到boundingbox标签
        xmin = int(float(xmlbox.find('xmin').text))
        ymin = int(float(xmlbox.find('ymin').text))
        xmax = int(float(xmlbox.find('xmax').text))
        ymax = int(float(xmlbox.find('ymax').text))

        #将boundingbox的四个坐标准换成中心点的坐标和boundingbox的宽高，并且转换映射到0-1
        cls_id = classes_name.index(cls_name)
        bb_h, bb_w = ymax - ymin, xmax - xmin
        if image_height == 0 or image_width == 0 :
            continue
        coodr_x, coord_y = (xmin + bb_w / 2) / image_width, (ymin + bb_h / 2) / image_height
        a_line = str(cls_id) + " " + str(coodr_x) + " " + str(coord_y) + " " + str(bb_w / image_width) + " " + str(
            bb_h / image_height) + "\n"
        all_lines.append(a_line)

    if have:                          #如果这张图片有我们需要的类，则保存这张图片,并写入到txt文件里
        txt_file = open(txt_save_path+file_name+".txt",'a')
        txt_file.writelines(all_lines)
        shutil.copy(img_path_origin+file_name+".jpg",img_savepath+file_name+".jpg")
        txt_file.close()
        count += 1
    xml_file.close()
print(f"---------转换了{count}个元素,结束-----------------")

